import torch
from models.fault_diagnosis import DiagnosisModel
from models.condition_diffusion_model import ConditionedDDPM

def test_model_dimensions():
    print("开始测试模型维度...")
    
    # 1. 测试特征提取器
    feature_extractor = DiagnosisModel(input_dim=4096, num_classes=5)
    test_input = torch.randn(8, 1, 4096)  # 批次大小为8
    
    # 测试特征提取
    features = feature_extractor.get_condition_features(test_input)
    print(f"特征提取器输出维度: {features.shape}")  # 应该是 [8, 64, 16]
    
    # 2. 测试DDPM模型
    ddpm = ConditionedDDPM(
        feature_extractor=feature_extractor,
        num_classes=5,
        cond_dim=256
    )
    
    # 3. 测试前向传播
    try:
        test_labels = torch.randint(0, 5, (8,))
        loss = ddpm(test_input, test_labels)
        print(f"DDPM训练损失: {loss.item()}")
        print("前向传播测试成功!")
    except Exception as e:
        print(f"前向传播测试失败: {str(e)}")
    
    # 4. 测试采样
    try:
        with torch.no_grad():
            # 获取条件特征
            cond_features = feature_extractor.get_condition_features(test_input)
            cond_features = cond_features.reshape(8, -1)  # 展平特征
            test_labels = torch.randint(0, 5, (8,), dtype=torch.long)  # 确保标签是 long 类型
            
            # 确保条件特征是浮点类型
            cond_features = cond_features.float()
            
            # 生成样本
            synthetic_signals, generated_labels = ddpm.sample(
                test_labels,  # 注意：这里调换了参数顺序，先传入labels
                cond_features,
                num_samples=1
            )
            print(f"生成信号维度: {synthetic_signals.shape}")  # 应该是 [8, 1, 4096]
            print(f"生成标签维度: {generated_labels.shape}")   # 应该是 [8]
            print("采样测试成功!")
    except Exception as e:
        print(f"采样测试失败: {str(e)}")

if __name__ == "__main__":
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    try:
        test_model_dimensions()
        print("\n所有测试完成!")
    except Exception as e:
        print(f"\n测试过程中出现错误: {str(e)}")